#!/usr/bin/python
#coding:utf-8

# Copyright 2011 Nicolau Leal Werneck, Anna Helena Reali Costa and
# Universidade de São Paulo
#
# Licensed under theglim Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################
## Open a (pinhole camera model) picture, find its orientation and
## extract the rectified.
##
## Changes that must be performed as soon as possible: read the
## intrinsic parameters from somewhere, and make it easy to switch to
## the equirectangular model.

import sys

import time

import matplotlib

if __name__ == '__main__':
    if sys.argv[0][-7:] == '-nox.py':
        matplotlib.use('Agg') 

import plot_aux

from pylab import *

import scipy.io

from scipy.special import dawsn
import scipy.ndimage

from numpy import dot

# from camori import PicturePinhole, PictureEqrec, quaternion_to_matrix, measure_error, quaternion_product, dir_colors
from camori_harris import PictureHarris

from quaternion import Quat, random_quaternion

import simplejson

import scipy.optimize
#from fmin_sa import Sampling
from simplex_sa import SimplexSO3

import Image

from plot_aux import *

set_printoptions(precision=7)

def all_oris(patch, a_i, sigma_i):
    Nfe = a_i.shape[0]
    Nf = 360*2
    deltatheta = 2*pi/(Nf)

    qqp = 1j*zeros(Nf)
    qqn = 1j*zeros(Nf)

    out = 1j*zeros(Nf)
    for m in range(Nfe):
        qqp[m] = (sigma_i[m] * a_i[m] * patch).sum()
        qqn[m] = ((-1)**m * sigma_i[m] * conj(a_i[m]) * patch).sum()

    for n in range(Nf):
        theta = n*deltatheta

        # qq = 1j*zeros((Nw,Nw))
        # for m in range(Nfe):
        #     qq += sigma_i[m] * a_i[m] * exp(1j * -m * theta)
        #     qq += (-1)**m * sigma_i[m] * conj(a_i[m]) * exp(1j * m * theta)
        # out[n] = (qq * patch).sum()

        for m in range(Nfe):
            out[n] += qqp[m] * exp(1j * -m * theta)
            out[n] += qqn[m] * exp(1j * m * theta)

    return out

def ori_target(theta,qqp,qqn):
    out = 1j * 0
    for m in range(Nfe):
        out += qqp[m] * exp(1j * -m * theta)
        out += qqn[m] * exp(1j * m * theta)
    return -(real(out)**2+imag(out)**2)

def find_ori(patch, a_i, sigma_i, tini=0.0):
    Nfe = a_i.shape[0]

    qqp = 1j*zeros(Nfe)
    qqn = 1j*zeros(Nfe)

    for m in range(Nfe):
        qqp[m] = (sigma_i[m] * a_i[m] * patch).sum()
        qqn[m] = ((-1)**m * sigma_i[m] * conj(a_i[m]) * patch).sum()
    
    ang = scipy.optimize.brent(ori_target, args=(qqp,qqn), brack=(tini-.01,tini+.01))
    
    return -ang

def find_ori_all(im, a_i, sigma_i):
    Nfe = a_i.shape[0]

    ang = zeros((im.shape[0], im.shape[1]))
    II = zeros((im.shape[0], im.shape[1]))

    sobel = array([[-1,0,1],[-2,0,2],[-1,0,1]])
    gx = scipy.ndimage.convolve(im, sobel)
    gy = scipy.ndimage.convolve(im, sobel.T)

    qqp = 1j*zeros((im.shape[0], im.shape[1], Nfe))
    qqn = 1j*zeros((im.shape[0], im.shape[1], Nfe))

    for m in range(Nfe):
        qqp[:,:,m] = scipy.ndimage.convolve(im, sigma_i[m] * a_i[m])
        qqn[:,:,m] = scipy.ndimage.convolve(im, (-1)**m * sigma_i[m] * conj(a_i[m]))

    # for j in range(im.shape[0]):
    #     for k in range(im.shape[1]):
    for j in range(100):
        for k in range(100):
            tini = arctan2(gy[j,k], gx[j,k])
            ang[j,k] = scipy.optimize.brent(ori_target, args=(qqp[j,k],qqn[j,k]), brack=(tini-.01,tini+.01))
            II[j,k] = -ori_target(ang[j,k],qqp[j,k],qqn[j,k])
            print j,k,ang[j,k], II[j,k]

    return ang,II



def myf(x):
    return 4.0/sqrt(pi) * (-x -dawsn(x) +2*x**2*dawsn(x))



if __name__ == '__main__' :

    rc('text',usetex=False)

    ## Avoid zero divide warnins...
    np.seterr(divide='ignore')

    PlotStuff=True
    if sys.argv[0][-7:] == '-nop.py':
        PlotStuff=False

    if PlotStuff:
        ## Plot stuff immediately
        ion()

    #################################################################
    ## Load image and initialize pic object

    ## Sets filename from input argument
    if len(sys.argv) < 3:
        print sys.argv[0], '<job_file.json> <frame_number>'
        raise Exception('Insufficient number of parameters')

    finput = open(sys.argv[1])
    job_params = simplejson.load(finput)
    finput.close()

    fileroot = job_params['root_directory']

    framenum = int(sys.argv[2])
    filename = fileroot+'/frames/'+job_params['filename_format']%framenum

    im = Image.open(filename)
    frame = array(im.convert('RGB'), dtype=float)
    imr = array(im.convert('RGB'), dtype=float)
    imr = imr[:,:,:3] #remove alpha channel

    im = imr.sum(2)/3

    Nfe = 12

    Nw = 11

    read_perona = np.load('perona_bases.npz')
    a_i = read_perona['a_i']
    sigma_i = read_perona['sigma_i']
    #a_i[0]-=mean(a_i[0])

    output = zeros(im.shape)

    sobel = array([[-1,0,1],[-2,0,2],[-1,0,1]])
    gx = scipy.ndimage.convolve(im, sobel)
    gy = scipy.ndimage.convolve(im, sobel.T)

    ang, II  = find_ori_all(im, a_i, sigma_i)

    figure(1)
    imshow(ang)
    figure(2)
    imshow(II)
    
